#!/usr/bin/python3
# filepath: /home/ma-user/work/OpenBOAT_zjh0/AscendBasicOps/AscendMatrixOps/Eye/EyeKernelLaunch_tiling/scripts/gen_data.py
import os
import numpy as np

def gen_eye_data():
    # 配置维度
    batch_size = 15
    rows = 777
    cols = 777
    shape = (batch_size, rows, cols)

    # 生成eye矩阵 (batch个相同的单位矩阵)
    single_eye = np.eye(rows, cols, dtype=np.float16)
    golden = np.repeat(single_eye[np.newaxis, ...], batch_size, axis=0)

    # 生成占位输入（若算子不需要真实输入，可用占位数据）
    input_x = np.zeros(shape, dtype=np.float16)

    # tiling数据
    tiling = np.array([batch_size, rows, cols], dtype=np.uint32)

    os.makedirs("input", exist_ok=True)
    os.makedirs("output", exist_ok=True)

    input_x.tofile("./input/input_x.bin")        # 占位
    tiling.tofile("./input/input_tiling.bin")
    golden.tofile("./output/golden.bin")

    print(f"Input (placeholder) shape: {input_x.shape}")
    print(f"Tiling: {tiling}")
    print(f"Golden (eye) shape: {golden.shape}")

if __name__ == "__main__":
    gen_eye_data()
